# analysis
rm(list=ls())
library(polywog)
library(MASS)
library(doParallel)
library(data.table)
library(ggplot2)
load("ensemble.Rdata")

dataset <- ensemble
dataset <- dataset[round==1076]
dataset$alpha1 <- ifelse(dataset$alpha==1, 1, 0)
dataset$alpha0 <- ifelse(dataset$alpha==0, 1, 0)
dataset$alpha5 <- ifelse(dataset$alpha>=.5,1,0)
dataset$sddistance0 <- ifelse(dataset$sd.distance==0, 1, 0)
dataset$discount0 <- ifelse(dataset$discount==0, 1, 0)
dataset$discount1 <- ifelse(dataset$discount==1, 1, 0)
dataset$meanAshare <- dataset$mean.Ashare
dataset$meanHGshare <- dataset$mean.HGshare
dataset$meanSGshare <- dataset$mean.SGshare
dataset$meanHshare <- dataset$mean.Hshare
dataset$sddistance <- dataset$sd.distance

polywog.select <- function(formula, data, degree=3, boot=5, select=.05, parallel=FALSE, lambda=NULL) {
  first <- polywog(formula=formula, data=data, degree=degree, boot=boot,
                   .parallel=parallel, nlambda=100, nfold=10, lambda=lambda)
  sig <- which((abs(summary(first)$coefficients[,1]/summary(first)$coefficients[,2])>1.96)==TRUE)
  rhs <- names(sig)
  rhs <- gsub("[.]", ")*I(", rhs) 
  if (rhs[1]=="(Intercept)") {
    rhs <- rhs[-1]
  } else {
    rhs[1] <- "-1"
  }
  if (length(rhs)>1) {
    rhs <- ifelse(rhs!="-1", paste("I(", rhs, ")", sep=""), "-1")
    rhs <- paste(rhs, collapse="+")
  } else if (length(rhs)==1) {
    rhs <- paste("I(", rhs, ")", sep="")
  } else if (length(rhs)==0){
    rhs <- "1"
  }  
  new.formula <- as.formula(paste(as.character(formula[2]), "~", rhs))
  final <- lm(new.formula, data)
}



predict.polywog <- function(pmodel, parties=3, discount=0, sddistance=0, alpha=0, memory=2, aspiration=.25, meanAshare=FALSE, 
                            meanSGshare=FALSE, meanHshare=FALSE, meanHGshare=FALSE) {
  #input.data <- expand.grid(0, 0, 0, parties, discount, sddistance, alpha, memory, aspiration, meanAshare, meanSGshare, meanHshare, meanHGshare, sharesticker)
  #names(input.data) <- c("eccentricity.mean", "representation.mean", "rep.gov.mean",
  if(length(parties)>1) {parties <- lapply(parties, function(x) expand.grid(x,seq(0,1,by=(1/x))))
                         parties <- do.call("rbind",parties)
                         #shares <- do.call("rbind",apply(parties, 1, functionX))
                         input.data <- expand.grid(0, parties[,1],discount, sddistance, alpha, memory, aspiration, 0,0,0,0)
                         input.data[,7 + which(c(meanAshare, meanSGshare, meanHshare, meanHGshare)==TRUE)] <- parties[,2]

  }
  if(length(parties)==1){  if(meanAshare==TRUE){input.data <- expand.grid(0, parties,discount, sddistance, alpha, memory, aspiration, seq(0,1,1/parties), 0, 0, 0)}
                           if(meanSGshare==TRUE){input.data <- expand.grid(0, parties,discount, sddistance, alpha, memory, aspiration, 0, seq(0,1,1/parties), 0, 0)}
                           if(meanHshare==TRUE){input.data <- expand.grid(0, parties,discount, sddistance, alpha, memory, aspiration, 0, 0, seq(0,1,1/parties), 0)}
                           if(meanHGshare==TRUE){input.data <- expand.grid(0, parties,discount, sddistance, alpha, memory, aspiration, 0, 0, 0, seq(0,1,1/parties))}}
  names(input.data) <- c(names(pmodel$model)[1],                      
                         "parties", "discount", "sddistance", "alpha", "memory", "aspiration", "meanAshare",
                         "meanSGshare", "meanHshare", "meanHGshare")
  names(input.data)[1] <- ifelse(names(input.data)[1]=="log(mean.gov.representation2)", "mean.gov.representation2", names(input.data)[1])
  input.data <- as.data.frame(input.data)
  output.data <- as.data.frame(input.data)
  input.data$alpha1 <- ifelse(input.data$alpha==1, 1, 0)
  input.data$alpha0 <- ifelse(input.data$alpha==0, 1, 0)
  input.data$discount1 <- ifelse(input.data$discount==1, 1, 0)
  input.data$discount0 <- ifelse(input.data$discount==0, 1, 0)
  input.data$sddistance0 <- ifelse(input.data$sddistance==0, 1, 0)
  input.data$alpha5 <- ifelse(input.data$alpha>=.5,1,0)
  input.data$memory2 <- ifelse(input.data$memory==2,1,0)
  input.data$memory10 <- ifelse(input.data$memory==10,1,0)
  input.data$aspiration25 <- ifelse(input.data$aspiration==0.25,1,0)
  input.data$aspiration9 <- ifelse(input.data$aspiration==0.9,1,0)
  input <- model.matrix(pmodel, data=input.data)
  betas <- mvrnorm(1000, coef(pmodel), vcov(pmodel))
  values <- input %*% t(betas)
  means <- apply(values, 1, mean)
  lower <- apply(values, 1, quantile, p=.025)
  upper <- apply(values, 1, quantile, p=.975)
  return(as.data.table(cbind(output.data, lower, means, upper,
                             type=paste0(as.character(pmodel$terms[[2]]), collapse=""))))
}




## cl <- makeCluster(4)
## registerDoParallel(cl)
## #Estimating Model Selections
## set.seed(64365485)
## model1_new<- polywog.select(mean.gov.representation2~parties+discount+sddistance+alpha+ aspiration + memory + meanAshare + meanHGshare + meanHshare + meanSGshare +
##                            alpha1+alpha0+sddistance0+discount0+discount1,
##                          data=dataset, degree=2, boot=40, parallel=TRUE)
## save(model1_new, file="analysis/regression.model1_new.RData")

## set.seed(64365485)
## model2 <- polywog.select(mean.representation~parties+discount+sddistance+alpha+ aspiration + memory + meanAshare + meanHGshare + meanHshare + meanSGshare +
##                            alpha1+alpha0+sddistance0+discount0+discount1,
##                          data=dataset, degree=2, boot=40, parallel=TRUE)
## save(model2, file="regression.model2.RData")

## set.seed(64365485)
## model3 <- polywog.select(mean.eccentricity~parties+discount+sddistance+alpha+ aspiration + memory + meanAshare + meanHGshare + meanHshare + meanSGshare +
##                            alpha1+alpha0+sddistance0+discount0+discount1,
##                          data=dataset, degree=2, boot=40, parallel=TRUE)
## save(model3, file="regression.model3.RData")
## stopCluster(cl)


library(ggplot2)
# Plotting predictions
plot.polywog <- function(x,parties, discount, sddistance, alpha, memory, aspiration, facets, plot=TRUE){
  pr.model2 <- predict.polywog(x,parties, discount, sddistance, alpha, memory, aspiration, TRUE, FALSE, FALSE, FALSE)
  pr.model3 <- predict.polywog(x,parties, discount, sddistance, alpha, memory, aspiration, FALSE, TRUE, FALSE, FALSE)
  pr.model4 <- predict.polywog(x,parties, discount, sddistance, alpha, memory, aspiration, FALSE, FALSE, TRUE, FALSE)
  pr.model5 <- predict.polywog(x,parties, discount, sddistance, alpha, memory, aspiration, FALSE, FALSE, FALSE, TRUE)
  
  means <- c(pr.model2$means, pr.model3$means, pr.model4$means, pr.model5$means)
  if(facets=="parties") { facet <- parties}
  if(facets=="discount") { facet <- discount}
  if(facets=="sddistance") { facet <- sddistance}
  if(facets=="alpha") { facet <- alpha}
  if(facets=="memory") { facet <- memory}
  if(facets=="aspiration") { facet <- aspiration}
  
  shares <- rep(pr.model2$meanAshare,4)
  names <- c(rep("Aggregator",dim(pr.model2)[1]),rep("Satisficing\nGovernator",dim(pr.model2)[1]),rep("Hunter",dim(pr.model2)[1]),rep("Governator",dim(pr.model2)[1]))
  facet <-  rep(as.matrix(pr.model2[,facets, with=FALSE]),4)
  
  upper <- c(pr.model2$upper, pr.model3$upper, pr.model4$upper, pr.model5$upper)
  lower <- c(pr.model2$lower, pr.model3$lower, pr.model4$lower, pr.model5$lower)
  data <- data.frame(means, names, shares, facet, upper, lower)
 if (plot==TRUE) {
  fig <- ggplot(data, aes(x=shares, y=means, group=names)) +
    facet_grid(~facet) +      
    geom_line(aes(colour=names))+
    geom_ribbon(aes(ymin=lower, ymax=upper), fill="grey73", alpha=0.2) +
    scale_x_continuous(breaks=c(0,0.5,1)) + 
    ylab("Predicted Difference to Sticker") + 
    theme_bw() +
    theme(legend.title=element_blank()) +
    theme(panel.grid.major.x = element_blank())+
    theme(panel.grid.major.y = element_blank())+
    theme(panel.grid.minor.x = element_blank())+
    theme(panel.grid.minor.y = element_blank())+
    theme(strip.background = element_blank())+
    scale_colour_manual("Rules",
                        values=c("red", "deeppink", "green", "blue"),
                        labels=c("Aggregator", "Governator", "Hunter",
                            "Satisficing\nGovernator"))
  return(fig)
 } else {
     return(data)
 }
}

load("analysis/regression.model1_new.Rdata")
load("analysis/regression.model2.Rdata")
load("analysis/regression.model3.Rdata")

plot.data <- plot.polywog(model1_new, parties=3:8, discount=.5, sddistance=.12,
                          alpha=0.5, memory=6, 
                     aspiration=.5, facets="parties", plot=FALSE)
plot.data <- as.data.table(plot.data)
plot.data$rule <- ifelse(plot.data$names=="Aggregator", "A", NA)
plot.data$rule <- ifelse(plot.data$names=="Governator", "G", plot.data$rule)
plot.data$rule <- ifelse(plot.data$names=="Hunter", "H", plot.data$rule)
plot.data$rule <- ifelse(plot.data$names=="Satisficing\nGovernator", "SG", plot.data$rule)
plot.data <- plot.data[,label.x:=1.1, list(names, facet)]
plot.data <- plot.data[,label.y:=means[shares==1], list(names, facet)]
plot.data$label.y <- ifelse(plot.data$facet==4&plot.data$rule=="H",
                            plot.data$label.y-.03, plot.data$label.y)
plot.data$label.y <- ifelse(plot.data$facet==5&plot.data$rule=="H",
                            plot.data$label.y+.03, plot.data$label.y)
fig1 <- ggplot(plot.data, aes(x=shares, y=means, group=names)) +
    facet_grid(~facet) +
    geom_ribbon(aes(ymin=lower, ymax=upper), fill="grey73", alpha=0.2) +    
    geom_line(aes(colour=names))+
    scale_x_continuous(breaks=c(0,0.5,1)) + 
    ylab("Government Misery:\nPredicted Difference to Sticker") + 
    theme_bw() +
    theme(legend.title=element_blank()) +
    theme(panel.grid.major.x = element_blank())+
    theme(panel.grid.major.y = element_blank())+
    theme(panel.grid.minor.x = element_blank())+
    theme(panel.grid.minor.y = element_blank())+
    theme(strip.background = element_blank())+
    theme(legend.position="none")+    
    geom_text(aes(x=label.x, y=label.y, label=rule))+
    theme(panel.margin = unit(.75, "lines"))+    
    xlab("Rule Shares of Non-Stickers")+    
    scale_colour_manual("Rules",
                        values=c("red", "deeppink", "green", "blue"),
                        labels=c("Aggregator", "Governator", "Hunter",
                            "Satisficing\nGovernator"))
fig1
ggsave(fig1, file="fig1.png", dpi=300)
plot.data1 <- plot.data
plot.data1$type <- 3

plot.data <- plot.polywog(model2, parties=3:8, discount=.5, sddistance=.12,
                          alpha=0.5, memory=6, 
                     aspiration=.5, facets="parties", plot=FALSE)
plot.data <- as.data.table(plot.data)
plot.data$rule <- ifelse(plot.data$names=="Aggregator", "A", NA)
plot.data$rule <- ifelse(plot.data$names=="Governator", "G", plot.data$rule)
plot.data$rule <- ifelse(plot.data$names=="Hunter", "H", plot.data$rule)
plot.data$rule <- ifelse(plot.data$names=="Satisficing\nGovernator", "SG", plot.data$rule)
plot.data$means <- plot.data$means*-1 #to change from representativeness to misery
plot.data$upper <- plot.data$upper*-1
plot.data$lower <- plot.data$lower*-1
plot.data <- plot.data[,label.x:=1.1, list(names, facet)]
plot.data <- plot.data[,label.y:=means[shares==1], list(names, facet)]
fig2 <- ggplot(plot.data, aes(x=shares, y=means, group=names)) +
    facet_grid(~facet) +
    geom_ribbon(aes(ymin=lower, ymax=upper), fill="grey73", alpha=0.2) +    
    geom_line(aes(colour=names))+
    scale_x_continuous(breaks=c(0,0.5,1)) + 
    ylab("Party-System Misery:\nPredicted Difference to Sticker") + 
    theme_bw() +
    theme(legend.title=element_blank()) +
    theme(panel.grid.major.x = element_blank())+
    theme(panel.grid.major.y = element_blank())+
    theme(panel.grid.minor.x = element_blank())+
    theme(panel.grid.minor.y = element_blank())+
    theme(strip.background = element_blank())+
    theme(legend.position="none")+
    theme(panel.margin = unit(.75, "lines"))+
    geom_text(aes(x=label.x, y=label.y, label=rule))+
    xlab("Rule Shares of Non-Stickers")+
    scale_colour_manual("Rules",
                        values=c("red", "deeppink", "green", "blue"),
                        labels=c("Aggregator", "Governator", "Hunter",
                            "Satisficing\nGovernator"))
fig2
ggsave(fig2, file="fig2.png", dpi=300)
plot.data2 <- plot.data
plot.data2$type <- 2

plot.data <- plot.polywog(model3, parties=seq(3,8,1), discount=.5, sddistance=mean(dataset$sddistance), alpha=0.5, memory=6, 
                     aspiration=.5, facets="parties", plot=FALSE)
plot.data <- as.data.table(plot.data)
plot.data$rule <- ifelse(plot.data$names=="Aggregator", "A", NA)
plot.data$rule <- ifelse(plot.data$names=="Governator", "G", plot.data$rule)
plot.data$rule <- ifelse(plot.data$names=="Hunter", "H", plot.data$rule)
plot.data$rule <- ifelse(plot.data$names=="Satisficing\nGovernator", "SG", plot.data$rule)
plot.data <- plot.data[,label.x:=1.1, list(names, facet)]
plot.data <- plot.data[,label.y:=means[shares==1], list(names, facet)]
fig3 <- ggplot(plot.data, aes(x=shares, y=means, group=names)) +
    facet_grid(~facet) +
    geom_ribbon(aes(ymin=lower, ymax=upper), fill="grey73", alpha=0.2) +
    geom_line(aes(colour=names))+
    scale_x_continuous(breaks=c(0,0.5,1)) + 
    ylab("Eccentricity:\nPredicted Difference to Sticker") + 
    theme_bw() +
    theme(legend.title=element_blank()) +
    theme(panel.grid.major.x = element_blank())+
    theme(panel.grid.major.y = element_blank())+
    theme(panel.grid.minor.x = element_blank())+
    theme(panel.grid.minor.y = element_blank())+
    theme(strip.background = element_blank())+
    theme(panel.margin = unit(.75, "lines"))+
    theme(legend.position="none")+
    geom_text(aes(x=label.x, y=label.y, label=rule))+    
    xlab("Rule Shares of Non-Stickers")+
    scale_colour_manual("Rules",
                        values=c("red", "deeppink", "green", "blue"),
                        labels=c("Aggregator", "Governator", "Hunter",
                            "Satisficing\nGovernator"))
fig3
ggsave(fig3, file="fig3.png", dpi=300)
